from tokenizers.implementations import ByteLevelBPETokenizer
import os
from transformers import PreTrainedTokenizerFast

from base_model.nlp.base_model.second_model_mla.UsherConfig import UsherConfig


def init():
    # 创建分词器实例
    tokenizer = ByteLevelBPETokenizer()
    config = UsherConfig()
    # 训练分词器
    tokenizer.train(
        files=[config.path + r'\train.jsonl'],
        vocab_size=config.vocab_size,
        min_frequency=2,
        special_tokens=["<｜pad｜>", "<｜unk｜>", "<｜begin_of_sentence｜>", "<｜end｜>", "<｜User｜>", "<｜Assistant｜>"]
    )

    # 新增路径创建逻辑
    if not os.path.exists(config.path):
        os.makedirs(config.path)

    tokenizer.save_model(config.path)

    # 创建Hugging Face的Tokenizer实例并保存配置
    hf_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        unk_token="<｜unk｜>",
        bos_token="<｜begin_of_sentence｜>",
        eos_token="<｜end｜>",
        pad_token="<｜pad｜>"
    )

    hf_tokenizer.pad_token_id = hf_tokenizer.convert_tokens_to_ids("<｜pad｜>")  # 获取实际ID
    # 保存完整配置
    hf_tokenizer.save_pretrained(config.path)  # 自动生成tokenizer_config.json
